#include "Tree_Of_Clusters.h"
#include <iostream>
#include <fstream>
#include <algorithm>
#include "Log.h"

template<class Funct>
void preorder(Cluster *root, Funct& functor) {
    if (root != NULL) {
        functor.apply(root);
        preorder(root->left,  functor);
        preorder(root->right, functor);
    }
}

template<class Funct>
void postrder(Cluster *root, Funct& functor) {
    if (root != NULL) {
        preorder(root->left,  functor);
        preorder(root->right, functor);
        functor.apply(root);
    }
}

Tree_Of_Clusters::Tree_Of_Clusters() {
    LOG_ENTER(TREE, "Tree_Of_Clusters::Tree_Of_Clusters");

    root = NULL;
    leaf_distances_stale = true;
    cluster_metric = Cluster_Metric(Cluster_Distance::COMPLETE,
                                    Representative_Set::WHOLE);

    LOG_EXIT();
}

// Destructor: traverse in preorder, deleting nodes
Tree_Of_Clusters::~Tree_Of_Clusters() {
    LOG_ENTER(TREE, "Tree_Of_Clusters::~Tree_Of_Clusters");

    destroy();

    LOG_EXIT();
}

class Node_Deleter {
public:
    void apply(Cluster *node) {
        delete node;
    }
};

void Tree_Of_Clusters::destroy() {
    Node_Deleter deleter;
    postorder(root, deleter);
}

// Copy constructor: traverse tree_of_clusters and build from root.
Tree_Of_Clusters::Tree_Of_Clusters(const Tree_Of_Clusters& other) {
    LOG_ENTER(TREE, "Tree_Of_Clusters::Tree_Of_Clusters");

    root = copy_tree(other.root);

    LOG_EXIT();
}

Cluster *Tree_Of_Clusters::copy_tree(const Cluster *node) {
    LOG_ENTER(TREE, "Tree_Of_Clusters::copy_tree");

    Cluster *result = NULL;
    if (node != NULL) {
        result = new CT_Node();
        result->left  = copy_tree(node->left);
        result->right = copy_tree(node->right);
        if (result->left != NULL) {
            result->left->parent = result;
        }
        if (result->right != NULL) {
            result->right->parent = result;
        }
    }

    LOG_EXIT();

    return result;
}

void Tree_Of_Clusters::set_root(Cluster *root_cluster) {
    root = root_cluster;
}

//
// Convenience function: make two children of a known parent.
//
void Tree_Of_Clusters::make_family(Cluster *parent,
                                   Cluster *A, Cluster *B) {
    LOG_ENTER(TREE, "Tree_Of_Clusters::make_family");
    LOG(TREE) << "A ID=" << A->ID << " B ID=" << B->ID << " parent ID=" << parent->ID;

    if (root == NULL) {
        set_root(parent);
    }

    parent->left  = A;
    parent->right = B;
    parent->antipoles.clear();

    LOG(TREE) << "parent @" << parent;
    LOG(TREE) << "A      @" << A;
    LOG(TREE) << "B      @" << B;

    leaf_distances_stale = true;

    LOG(TREE) << "New tree_of_clusters:";
    preorder_write(Log::output);

    LOG_EXIT();
}

// Connect two children to a known parent node.
// Children may not have tree nodes yet.  If so, create them.
void Tree_Of_Clusters::set_children(Cluster *parent,
                                    Cluster *L, Cluster *R) {
    LOG_ENTER(TREE, "Tree_Of_Clusters::set_children");
    LOG(TREE) << "L ID=" << L->ID << " R ID=" << R->ID << " parent ID=" << parent->ID;

    parent->left  = L;
    parent->right = R;
    parent->antipoles.clear();

    LOG(TREE) << "parent @" << parent;
    LOG(TREE) << "L      @" << L;
    LOG(TREE) << "R      @" << R;

    leaf_distances_stale = true;

    LOG(TREE) << "New tree_of_clusters:";
    preorder_write(Log::output);

    LOG_EXIT();
}

// Connect two existing children to their parent.
// Parent may not have a tree node yet.  If not, create it.
void Tree_Of_Clusters::set_parent(Cluster *parent,
                                  Cluster *L, Cluster *R) {
    LOG_ENTER(TREE, "Tree_Of_Clusters::set_parent");
    LOG(TREE) << "L ID=" << L->ID << " R ID=" << R->ID << " parent ID=" << parent->ID;

    L->parent = parent;
    R->parent = parent;

    LOG(TREE) << "parent @" << parent;
    LOG(TREE) << "L      @" << L;
    LOG(TREE) << "R      @" << R;

    leaf_distances_stale = true;

    LOG(TREE) << "New tree_of_clusters:";
    preorder_write(Log::output);

    LOG_EXIT();
}

void Tree_Of_Clusters::add_vertex_to_graph(Cluster *n, ct_graph_t& graph) const {
    LOG_ENTER(TREE, "Tree_Of_Clusters::add_vertex_to_graph");
    LOG(TREE) << "n = @" << n;

    if (n != NULL) {
        LOG(TREE) << "n->ID=" << n->ID;

        CT_Vertex_Info vInfo(n->ID, n->max_leaf_distance);
        graph.add_vertex(vInfo, ID);
        add_vertex_to_graph(n->left,  graph);
        add_vertex_to_graph(n->right, graph);
    }

    LOG_EXIT();
}

// Connect the vertices of the graph, following the convention
// that each vertex has exactly three neighbors, in the following
// order: parent, left, right.
void Tree_Of_Clusters::add_edges_to_graph(int parent_ID,
                                          Cluster *n, ct_graph_t& graph) const {
    LOG_ENTER(TREE, "Tree_Of_Clusters::add_edges_to_graph");
    LOG(TREE) << "n = @" << n;

    if (n != NULL) {
        LOG(TREE) << "n->ID=" << n->ID;

        int ID       = n->ID;
        int left_ID  = n->left  == NULL ? -1 : n->left-> cluster->ID;
        int right_ID = n->right == NULL ? -1 : n->right->cluster->ID;
        graph.add_edge(ID, parent_ID, CT_Edge_Info(0));
        graph.add_edge(ID, left_ID,   CT_Edge_Info(0));
        graph.add_edge(ID, right_ID,  CT_Edge_Info(0));

        add_edges_to_graph(ID, n->left,  graph);
        add_edges_to_graph(ID, n->right, graph);
    }

    LOG_EXIT();
}

ct_graph_t* Tree_Of_Clusters::to_graph() {
    LOG_ENTER(TREE, "Tree_Of_Clusters::to_graph");

    set_leaf_distances();

    ct_graph_t *graph = new ct_graph_t();

    // convert the tree_of_clusters into a graph.  Each vertex will have,
    // at most, three neighbors, in the following order
    // (order matters)
    // first:  parent
    // second: left child
    // third:  right child
    // The root node, however, has no parent, but has two neighbors
    // (its children), so we need a special vertex with ID = -1 to denote
    // an absent neighbor.  Hence, leaf nodes will be connected to this
    // special vertex, and so will the root.

    // Here's the special (ID = -1) vertex
    graph->add_vertex(CT_Vertex_Info(-1, -1), -1);

    // now recurse down the tree_of_clusters, adding vertices in preorder
    add_vertex_to_graph(root, *graph);

    // and finally connect the added vertices.
    add_edges_to_graph(-1, root, *graph);

    LOG_EXIT();

    return graph;
}

void Tree_Of_Clusters::init_from_graph(ct_graph_t& graph,
                                       std::vector<Cluster*>& clusters) {
    LOG_ENTER(TREE, "Tree_Of_Clusters::init_from_graph");

    // zap the old tree_of_clusters, reclaiming memory
    delete_tree(root);

    boost::unordered_map<int, Cluster*> ID_to_cluster;
    ID_to_node[-1] = NULL;

    // Build a map of ID to Cluster*
    for (std::vector<Cluster*>::iterator iter = clusters.begin();
         iter != clusters.end();
         iter++) {
        int ID = *iter->ID;
        ID_to_cluster[ID] = *iter;
    }

    ct_vertex_range_t vertices = graph.get_vertices();
    ct_vertex_iter_t viter = vertices.first;
    ct_vertex_iter_t vend  = vertices.second;

    // now go through the vertices again, connecting the nodes
    // we just created above.
    viter = vertices.first;  // rewind the iterator for another pass
    int root_ID = -1;
    for (; viter != vend; viter++) {
        ct_adjacency_vertex_range_t neighbors
            = graph.get_adjacent_vertices(*viter);
        // iterate through the three neighbors.
        ct_adjacency_iter_t niter = neighbors.first;
        ct_adjacency_iter_t nend  = neighbors.second;

        // first neighbor is the parent
        int parent_ID = graph.properties(*niter++).cluster_ID;

        // second neighbor is left child
        int left_ID = -1;
        if (niter != nend)
            left_ID = graph.properties(*niter++).cluster_ID;

        // and right neighbor is right child
        int right_ID = -1;
        if (niter != nend)
            right_ID = graph.properties(*niter++).cluster_ID;

        // now connect everyone
        CT_Vertex_Info vinfo = graph.properties(*viter);
        int ID = vinfo.cluster_ID;
        Cluster *node = ID_to_cluster[ID];
        node->parent = ID_to_node[parent_ID];
        node->left   = ID_to_node[left_ID];
        node->right  = ID_to_node[right_ID];

        if (parent_ID == -1)
            root = node;
    }

    LOG_EXIT();
}

void Tree_Of_Clusters::write_nodes(std::string& file_name) {
    LOG_ENTER(TREE, "Tree_Of_Clusters::write_nodes");

    std::ofstream file;
    file.open(file_name.c_str());

    std::unordered_map<int, Cluster*>::iterator iter;
    std::vector<int> IDs;
    for (iter = ID_to_node.begin(); iter != ID_to_node.end(); iter++) {
        IDs.push_back((*iter).first);
    }
    std::sort(IDs.begin(), IDs.end());

    LOG(TREE) << "sorted the IDs";

    for (size_t i=0; i<IDs.size(); i++) {

        LOG(TREE) << "IDs[" << i << "]=" << IDs.at(i);

        int ID = IDs.at(i);
        file << ID << " ";

        Cluster* node = ID_to_node[ID];
        node->items->write_IDs(file);
        file << "\n";
    }

    LOG_EXIT();
}


void Tree_Of_Clusters::write(std::string& file_name) {
    LOG_ENTER(TREE, "Tree_Of_Clusters::write");
    LOG(TREE) << " file=" << file_name << "\n";

    std::ofstream file;
    file.open(file_name.c_str());

    write(file);

    file.close();

    LOG_EXIT();
}

void Tree_Of_Clusters::write(std::ostream& os) {
    LOG_ENTER(TREE, "Tree_Of_Clusters::write(stream)");

    ct_graph_t *graph = to_graph();

    os << *graph;
    delete graph;

    LOG_EXIT();
}

void Tree_Of_Clusters::preorder_write(std::ostream& os) {
    preorder_write(os, root);
}

class Node_Writer {
public:
    std::ostream *os;
    Node_Writer(std::ostream& os) {
        this->os = &os;
    }

    void apply(Cluster *node) {
        int cluster_ID = -1;
        int parent_ID  = -1;
        int left_ID    = -1;
        int right_ID   = -1;
        if (node != NULL) cluster_ID = node->ID;

        if (node->parent != NULL &&
            node->parent != NULL) parent_ID = node->parent->ID;

        if (node->left   != NULL &&
            node->left-  != NULL) left_ID   = node->left->ID;

        if (node->right  != NULL &&
            node->right  != NULL) right_ID  = node->right->ID;

        (*os) << "Node (cluster=" << cluster_ID
              << " left=" << left_ID
              << " right=" << right_ID
              << " parent=" << parent_ID << ") @" << node << "\n";
    }
};

void Tree_Of_Clusters::preorder_write(std::ostream& os, Cluster *node) {
    Node_Writer writer(os);
    preorder(root, writer);
}

void Tree_Of_Clusters::set_leaf_distances() {
    LOG_ENTER(TREE,"Tree_Of_Clusters::set_leaf_distances");
    LOG(TREE) << "distances stale=" << (leaf_distances_stale ? "YES" : "NO");

    if (leaf_distances_stale) {
        set_leaf_distances(root);
        leaf_distances_stale = false;
    }

    LOG_EXIT();
}

class Leaf_Distance_Setter {
public:
    Cluster_Metric metric;

    Leaf_Distance_Setter(Cluster_Metric metric) {
        this->metric = metric;
    }

    void apply(Cluster *node) {
        if (node->is_leaf())
            node->max_leaf_distance = 0;
        else
            node->max_leaf_distance =
                cluster_metric.distance(node->left, node->right);
    }
};

void Tree_Of_Clusters::set_leaf_distances(CT_Node *node) {
    LOG_ENTER(TREE,"Tree_Of_Clusters::set_leaf_distances(node)");
    LOG(TREE) << "node = @" << node;

    Leaf_Distance_Setter leafer(cluster_metric);
    preorder(root, leafer);

    LOG_EXIT();
}

// Modify the tree_of_clusters, replacing a node for a given old cluster
// with a whole new subtree.
// If the new node is NULL, we're effectively deleting the old cluster
void Tree_Of_Clusters::replace_node(Cluster *old_cluster,
                                    Cluster *new_cluster) {
    LOG_ENTER(TREE,"Tree_Of_Clusters::replace_node");

    LOG(TREE) << "old_cluster ID=" << old_cluster->ID;
    if (new_cluster == NULL)
        LOG(TREE) << "new_cluster is NULL";
    else
        LOG(TREE) << "new_cluster ID=" << new_cluster->ID;

    if (old_cluster == NULL) {
        LOG(TREE) << "old_cluster is NULL, skipping replacement";
    }
    else {
        Cluster *parent = old_cluster->parent;
        if (parent == NULL) {
            LOG(ERROR) << "null parent in replace_node (old node ID = "
                       << old_cluster->ID;
        }
        else {
            bool is_left_child;
            is_left_child = (old_cluster == parent->left);

            LOG(TREE) << "parent ID=" << parent->ID
                      << " is_left_child=" << (is_left_child ? "true" : "false");

            // zap node and children.
            delete_tree(old_cluster);

            // splice in new node (may be NULL, that's OK)
            if (is_left_child)
                parent->left = new_cluster;
            else
                parent->right = new_cluster;
            new_cluster->parent = parent;
        }
    }

    LOG(TREE) << "spliced child in.";
    LOG_EXIT();
}

void Tree_Of_Clusters::replace_cluster(Cluster *old_cluster,
                                       Tree_Of_Clusters *new_subtree) {
    // Find old node, its parent, and its child relationship
    LOG_ENTER(TREE,"Tree_Of_Clusters::replace_cluster (with subtree)");

    replace_node(old_cluster, new_subtree->root);

    LOG_EXIT();
}

void Tree_Of_Clusters::delete_cluster(Cluster *victim) {
    replace_cluster(victim, (Cluster*)NULL);
}

void Tree_Of_Clusters::push_clusters(Cluster *L, Cluster *R) {
    pending_subtrees.push(L);
    pending_subtrees.push(R);
}

void Tree_Of_Clusters::pop_clusters(Cluster*& L, Cluster*& R) {
    R = pending_clusters.top();
    pending_clusters.pop();

    L = pending_clusters.top();
    pending_clusters.pop();
}

class Node_Pusher {
public:
    std::stack<Cluster*> *pending_clusters;
    Node_Pusher(std::stack<Cluster*>& pending_clusters) {
        this->pending_clusters = &pending_clusters;
    }

    void action(Cluster *node) {
        pending_clusters->push(node->left);
        pending_clusters->push(node->right);
    }
};
void Tree_Of_Clusters::push_tree_nodes(Tree_Of_Clusters& subtree) {
    LOG_ENTER(TREE, "Tree_Of_Clusters::push_tree_nodes");

    Node_Pusher pusher(pending_subtrees);
    preorder(subtree.root, pusher);

    LOG_EXIT();
}

